#include "USBHub.h"
#include "USBPort.h"

#define BUFFER_SIZE 2048

#pragma warning( disable : 4201 )

typedef union _USB_PROTOCOLS {
	ULONG  ul;

	struct {
		ULONG Usb110 :1;
		ULONG Usb200 :1;
		ULONG Usb300 :1;
		ULONG ReservedMBZ  :29;
	};

} USB_PROTOCOLS, *PUSB_PROTOCOLS;

typedef union _USB_NODE_CONNECTION_INFORMATION_EX_V2_FLAGS {
	ULONG  ul;

	struct {
		ULONG DeviceIsOperatingAtSuperSpeedOrHigher  :1;
		ULONG DeviceIsSuperSpeedCapableOrHigher  :1;
		ULONG DeviceIsOperatingAtSuperSpeedPlusOrHigher  :1;
		ULONG DeviceIsSuperSpeedPlusCapableOrHigher  :1;
		ULONG ReservedMBZ  :28;
	};

} USB_NODE_CONNECTION_INFORMATION_EX_V2_FLAGS, *PUSB_NODE_CONNECTION_INFORMATION_EX_V2_FLAGS;

#define IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX_V2  0x22045C

#pragma warning( default : 4201 )


typedef struct _USB_NODE_CONNECTION_INFORMATION_EX_V2 {
	// one based port number
	ULONG  ConnectionIndex;

	// length of the structure
	ULONG  Length;

	// On input a bitmask that indicates which USB protocols are understood by the caller
	// On output a bitmask that indicates which USB signaling protocols are supported by the port
	USB_PROTOCOLS SupportedUsbProtocols;

	// A bitmask indicating properties of the connected device or port
	USB_NODE_CONNECTION_INFORMATION_EX_V2_FLAGS Flags;
} USB_NODE_CONNECTION_INFORMATION_EX_V2, *PUSB_NODE_CONNECTION_INFORMATION_EX_V2;

USBHub::USBHub(void)
: m_hubPortCount(0), m_hubIsBusPowered(false), m_hubIsRootHub(false)
{
}

USBHub::~USBHub(void)
{
	FreeUSBPorts();
}

const std::vector<USBPort*>* USBHub::GetPorts()
{
	FreeUSBPorts();

	HANDLE hHub = CreateFile(m_hubDevicePath.c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
	if (hHub != INVALID_HANDLE_VALUE)
	{
		BYTE *buffer = new BYTE[BUFFER_SIZE];
		if (buffer)
		{
			for (int i = 1; i <= m_hubPortCount; ++i)
			{
				USB_NODE_CONNECTION_INFORMATION_EX *pConnectionInfo = (USB_NODE_CONNECTION_INFORMATION_EX*)buffer;
				pConnectionInfo->ConnectionIndex = i;
				DWORD bytesReturned;
				if (DeviceIoControl(hHub, IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX, pConnectionInfo, BUFFER_SIZE, pConnectionInfo, BUFFER_SIZE, &bytesReturned, NULL))
				{
					USBPort *pUSBPort = new USBPort();
					pUSBPort->SetPortNumber(i);
					pUSBPort->SetHubDevicePath(m_hubDevicePath.c_str());
					pUSBPort->SetStatus(pConnectionInfo->ConnectionStatus);
					pUSBPort->SetSpeed((USB_DEVICE_SPEED)pConnectionInfo->Speed);
					pUSBPort->SetDeviceConnected(pConnectionInfo->ConnectionStatus == DeviceConnected);
					pUSBPort->SetDeviceIsHub(pConnectionInfo->DeviceIsHub ? true : false);
					pUSBPort->SetDeviceDescriptor(pConnectionInfo->DeviceDescriptor);

					USB_NODE_CONNECTION_INFORMATION_EX_V2 connectionInfoV2;
					connectionInfoV2.ConnectionIndex = i;
					connectionInfoV2.Length = sizeof(connectionInfoV2);
					connectionInfoV2.SupportedUsbProtocols.Usb300 = 1;
					if (DeviceIoControl(hHub, IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX_V2, &connectionInfoV2, sizeof(connectionInfoV2), &connectionInfoV2, sizeof(connectionInfoV2), &bytesReturned, NULL))
					{
						if (connectionInfoV2.Flags.DeviceIsOperatingAtSuperSpeedOrHigher)
						{
							pUSBPort->SetSpeed((USB_DEVICE_SPEED)USBSuperSpeed);
						}

						if (connectionInfoV2.Flags.DeviceIsOperatingAtSuperSpeedPlusOrHigher)
						{
							pUSBPort->SetSpeed((USB_DEVICE_SPEED)USBSuperSpeedPlus);
						}
					}

					m_ports.push_back(pUSBPort);
				}
			}

			delete [] buffer;
		}

		CloseHandle(hHub);
	}

	return &m_ports;
}

void USBHub::FreeUSBPorts()
{
	for (std::vector<USBPort*>::iterator	iter = m_ports.begin();
											iter != m_ports.end();
											++iter)
	{
		USBPort *pUSBPort = *iter;
		delete pUSBPort;
	}

	m_ports.clear();
}